package jtg.generator;

import com.microsoft.z3.*;
import com.microsoft.z3.Context;
import javafx.util.Pair;
import jtg.graphics.SootCFG;
import jtg.parser.Z3Parser;
import jtg.state.BranchPathState;
import jtg.visualizer.Visualizer;
import soot.*;
import soot.jimple.internal.*;
import soot.toolkits.graph.UnitGraph;

import java.io.File;
import java.util.*;

public class DataFlowGenerator extends Generator{

    Map<String, Unit> defNodeMap = new HashMap<>();
    Map<String, Unit> useNodeMap = new HashMap<>();

    Map<Pair<Unit,Unit>,Integer> duPairMap = new HashMap<>();

    Map<String,Set<Unit>> defNodeSetMap = new HashMap<>();
    Map<String,Set<Unit>> useNodeSetMap = new HashMap<>();


    int branchNum = 0;

    public DataFlowGenerator(String className, String methodName) {
        super(className, methodName);
        String defaultClsPath = System.getProperty("user.dir") + File.separator + "target" + File.separator + "classes";
        new PrimeCoverageGenerator(defaultClsPath, className, methodName);
    }

    public DataFlowGenerator(String classPath, String className, String methodName) {
        super(classPath, className, methodName);
    }

    @Override
    void drawCFG() {
        drawCFG(clsName + "." + mtdName, false);
    }

    public void drawCFG(String graphName, boolean indexLabel) {
        Visualizer.printCFGDot(graphName, ug, indexLabel);
    }

    private List<Local> getParameters() {
        return new ArrayList<>(body.getParameterLocals());
    }

    public List<String> generate() {

        ArrayList<String> testSet = null;

        System.out.println("============================================================================");
        System.out.println("Generating test case inputs for method: " + clsName + "." + mtdName + "()");
        System.out.println("============================================================================");
        try {
            testSet = new ArrayList<String>();
            Context context = new Context();
            Solver solver = context.mkSolver();

            if(ug.getHeads().size()>1) throw new Exception("The number of heads is more than 1");
            Unit h = ug.getHeads().get(0);
            // path = ug.getExtendedBasicBlockPathBetween(h, t);
            for (Unit unit: ug) {
                if(unit instanceof JIfStmt){
                    branchNum++;
                }
            }
            List<BranchPathState> pathStates = initPathState(context, h);
            calculateWithBFS(context, ug, new HashSet<Unit>(), pathStates, testSet);
            for(Pair<Unit,Unit> pair : duPairMap.keySet()){
                System.out.println("The def-use pair is: " + pair);
            }
            int allPathNum = calculateAllPathNum();
            int visitedPathNum = duPairMap.size();
            System.out.println("The number of all paths is: " + allPathNum+ " The number of visited paths is: " + visitedPathNum);



        } catch (Exception e) {
            System.err.println("Error in generating test cases: ");
            System.err.println(e.toString());
            e.printStackTrace();
        }
        if (!testSet.isEmpty()) {
            System.out.println("");
            System.out.println("The generated test case inputs:");
            int count = 1;
            for (String tc : testSet) {
                System.out.println("( " + count + " ) " + tc.toString());
                count++;
            }
        }
        return testSet;
    }

    public List<BranchPathState> initPathState(Context context, Unit h){
        List<BranchPathState> pathStates = new ArrayList<>();
        // 创建初始状态
        BranchPathState initialState = new BranchPathState(h);
        for(Local local : getParameters()){
            String parName = local.getName();
            initialState.params.put(parName, context.mkIntConst(parName));
            initialState.paramLatestMap.put(parName,0);
        }
        initialState.unitPath.add(h);
        pathStates.add(initialState);
        return pathStates;
    }

    public int calculateAllPathNum(){
        int sum = 0;
        for(String str : defNodeSetMap.keySet()){
            if(useNodeSetMap.containsKey(str)){
                sum += defNodeSetMap.get(str).size()*useNodeSetMap.get(str).size();
            }

        }
        return sum;
    }

    public void calculateWithBFS(Context context, UnitGraph ug, Set<Unit> visitedNode, List<BranchPathState> pathStates, List<String> testSet) throws Exception {
        while(!pathStates.isEmpty()){
            BranchPathState pathState = pathStates.remove(0);
            Unit unit = pathState.node;
            if(pathState.unitPath.size()>50||unit instanceof JReturnStmt ||unit instanceof JReturnVoidStmt){
                if(pathCheck(pathState.unitPath)){
                    List<BoolExpr> pathConstraint = calPathConstraint(context, pathState.unitPath);
                    String testSetByPathConstraint = getTestSetByPathConstraint(context, pathConstraint);
                    if(testSetByPathConstraint!=null){
                        testSet.add(testSetByPathConstraint);
                    }
                }
                continue;
            }

            Solver solver = context.mkSolver();
            List<BoolExpr> pathConstraint1 = calPathConstraint(context, pathState.unitPath);
            if(!pathConstraint1.isEmpty()) {
                for(BoolExpr expr : pathConstraint1 ){
                    solver.add(expr);
                }

                if(solver.check() == Status.UNSATISFIABLE){
                    continue;
                }
            }


            visitedNode.add(unit);
            if(unit instanceof JIfStmt){
                // 如果是if语句，那么就是两条路径
                // 生成两个新的状态
                BranchPathState pathState1 = new BranchPathState(pathState);
                BranchPathState pathState2 = new BranchPathState(pathState);
                // 生成两个新的约束
                // 有可能是一个循环，判断一下

                // BoolExpr ifstms = (BoolExpr) getZ3ExprFromSootExpr(context, pathState.params, ((JIfStmt) unit).getCondition());
                List<Unit> nextUnitList = ug.getSuccsOf(unit);
                Unit nextUnit1 = nextUnitList.get(0);
                Unit nextUnit2 = nextUnitList.get(1);
                pathState1.node = nextUnit1;
                pathState2.node = nextUnit2;
                pathState1.unitPath.add(nextUnit1);
                pathState2.unitPath.add(nextUnit2);
                // pathState1.PathConstraints.add(context.mkNot(ifstms));
                // pathState2.PathConstraints.add(ifstms);


                pathStates.add(pathState1);
                pathStates.add(pathState2);
            }else if(unit instanceof JAssignStmt){
                // 如果是赋值语句，那么就是一条路径
                // 生成一个新的状态
                BranchPathState pathState1 = new BranchPathState(pathState);
                // 生成一个新的约束
                String leftOp = ((JAssignStmt) unit).getLeftOp().toString();
                pathState1.params.put(leftOp, context.mkIntConst(leftOp));
                // Expr rightOp = getZ3ExprFromSootExpr(context, pathState1.params, ((JAssignStmt) unit).getRightOp());
                // if (rightOp == null) {
                //    System.out.println("Ignore the statement: " + unit.toString());
                //    continue;
                //}
                // pathState1.PathConstraints.add(context.mkEq(pathState1.params.get(leftOp), rightOp));
                pathState1.node = ug.getSuccsOf(unit).get(0);
                pathState1.unitPath.add(pathState1.node);
                pathStates.add(pathState1);

            }else if(unit instanceof JGotoStmt){
                // 如果是goto语句，那么就是一条路径
                // 生成一个新的状态
                BranchPathState pathState1 = new BranchPathState(pathState);
                // 判断是否是循环
                Unit nextUnit = ug.getSuccsOf(unit).get(0);

                pathState1.node = nextUnit;
                pathState1.unitPath.add(pathState1.node);

                pathStates.add(pathState1);
            }else {
                // 如果是其它语句，那么就是一条路径
                // 生成一个新的状态
                BranchPathState pathState1 = new BranchPathState(pathState);
                // 生成一个新的约束
                pathState1.node = ug.getSuccsOf(unit).get(0);
                pathState1.unitPath.add(pathState1.node);
                pathStates.add(pathState1);
            }
        }
    }

    private boolean pathCheck(List<Unit> unitPath) {
        boolean flag = false;
        for (int i = 0; i < unitPath.size(); i++) {
            Unit unit = unitPath.get(i);
            if(unit instanceof JAssignStmt||unit instanceof JIdentityStmt){
                String def = null;
//                List<ValueBox> defBoxes = unit.getDefBoxes();
//                System.out.println("The defBoxes is: " + defBoxes);
//                for(ValueBox valueBox : defBoxes){
//                    if(valueBox.getValue() instanceof Local)
//                        def = ((Local) valueBox.getValue()).getName();
//                    if(defNodeSetMap.containsKey(def)) {
//                        defNodeSetMap.get(def).add(unit);
//                    }else {
//                        Set<Unit> set = new HashSet<>();
//                        set.add(unit);
//                        defNodeSetMap.put(def,set);
//                    }
//                    for(int j = i+1; j < unitPath.size(); j++){
//                        String use = null;
//                        Unit unit1 = unitPath.get(j);
//                        List<ValueBox> useBoxes = unit1.getUseBoxes();
//                        for(ValueBox valueBox1 : useBoxes){
//                            if(valueBox1.getValue() instanceof Local)
//                                use = ((Local) valueBox1.getValue()).getName();
//                            if(def.equals(use)){
//                                if(duPairMap.containsKey(new Pair<>(unit,unit1))){ continue; }
//                                flag = true;
//                                duPairMap.put(new Pair<>(unit,unit1),1);
//                                if(useNodeSetMap.containsKey(def)) {
//                                    useNodeSetMap.get(def).add(unit1);
//                                }else {
//                                    Set<Unit> set = new HashSet<>();
//                                    set.add(unit1);
//                                    useNodeSetMap.put(def,set);
//                                }
//                            }
//                        }
//                    }
//
//                }
                if(unit instanceof JAssignStmt){
                    def = ((JAssignStmt) unit).getLeftOp().toString();
                    if(defNodeSetMap.containsKey(def)) {
                        defNodeSetMap.get(def).add(unit);
                    }else {
                        Set<Unit> set = new HashSet<>();
                        set.add(unit);
                        defNodeSetMap.put(def,set);
                    }
                }else{
                    def = ((JIdentityStmt) unit).getLeftOp().toString();
                    if(defNodeSetMap.containsKey(def)) {
                        defNodeSetMap.get(def).add(unit);
                    }else {
                        Set<Unit> set = new HashSet<>();
                        set.add(unit);
                        defNodeSetMap.put(def,set);
                    }
                }
                for(int j = i+1; j < unitPath.size(); j++){
                    Unit unit1 = unitPath.get(j);
                    if(unit1 instanceof JAssignStmt){
                        String rightOp = ((JAssignStmt) unit1).getRightOp().toString();
                        if(rightOp.contains(def+" ")||rightOp.contains(" "+def)){
                            if(duPairMap.containsKey(new Pair<>(unit,unit1))){ continue; }
                            flag = true;
                            duPairMap.put(new Pair<>(unit,unit1),1);
                            if(useNodeSetMap.containsKey(def)) {
                                useNodeSetMap.get(def).add(unit1);
                            }else {
                                Set<Unit> set = new HashSet<>();
                                set.add(unit1);
                                useNodeSetMap.put(def,set);
                            }
                        }
                        String leftOp1 = ((JAssignStmt) unit1).getLeftOp().toString();
                        if(def.equals(leftOp1)){
                            break;
                        }
                    }else{
                        String op = unit1.toString();
                        if(op.contains(def+" ")||op.contains(" "+def)){
                            if(duPairMap.containsKey(new Pair<>(unit,unit1))){ continue; }
                            flag = true;
                            if(useNodeSetMap.containsKey(def)) {
                                useNodeSetMap.get(def).add(unit1);
                            }else {
                                Set<Unit> set = new HashSet<>();
                                set.add(unit1);
                                useNodeSetMap.put(def,set);
                            }
                            duPairMap.put(new Pair<>(unit,unit1),1);
                        }
                    }
                }
            }
        }
        return flag;
    }


    public List<BoolExpr> calPathConstraint(Context context, List<Unit> path) {
        Map<String, Expr> paras = new HashMap<>();
        Map<String, Integer> latestAssign = new HashMap<>();
        Z3Parser parser = new Z3Parser(context, paras, latestAssign);
        List<BoolExpr> pathConstraint = new ArrayList<>();
        for (Local local : getParameters()) {
            latestAssign.put(local.getName(), 0);
            paras.put(local.getName() + "_0", parser.fromLocal(local, pathConstraint));
        }
        String expectedResult = "";

        for (int i = 0; i < path.size(); i++) {
            Unit stmt = path.get(i);

            if (stmt instanceof JAssignStmt) {
                // assignList.put(((JAssignStmt) stmt).getLeftOp().toString(), ((JAssignStmt) stmt).getRightOp().toString());
                // String leftOp = ((JAssignStmt) stmt).getLeftOp().toString();
                // paras.put(leftOp, context.mkIntConst(leftOp));
                Expr rightOp = parser.fromSootExpr(((JAssignStmt) stmt).getRightOp());
                if (rightOp == null) {
                    // leftOp 设置为对应类型任意值
                    parser.assignLocal((Local) ((JAssignStmt) stmt).getLeftOp(), pathConstraint);
                    System.err.println("Ignore the statement: " + stmt + " assign arbitrary value to " + ((JAssignStmt) stmt).getLeftOp().toString());
                    continue;
                }
                Value leftOp = ((JAssignStmt) stmt).getLeftOp();
                if (leftOp instanceof Local) {
                    Expr newLeft = parser.assignLocal((Local) leftOp, pathConstraint);
                    pathConstraint.add(context.mkEq(newLeft, rightOp));
                } else if (leftOp instanceof JArrayRef) {
                    Value base = ((JArrayRef) leftOp).getBase();
                    Value index = ((JArrayRef) leftOp).getIndex();
                    if (base instanceof Local) {
                        ArrayExpr rightStore = context.mkStore(parser.fromLocal((Local) base, pathConstraint), parser.fromSootExpr(index), rightOp);
                        ArrayExpr newLeft = (ArrayExpr) parser.assignLocal((Local) base, pathConstraint);
                        pathConstraint.add(context.mkEq(newLeft, rightStore));
                    } else {
                        System.out.println("Ignore the statement: " + stmt.toString());
                        continue;
                    }
                } else {
                    System.out.println("Ignore the statement: " + stmt.toString());
                    continue;
                }
                continue;
            }
            if (stmt instanceof JIfStmt) {
//                System.out.println("The ifstmt is: " + stmt);
                BoolExpr ifstms = (BoolExpr) parser.fromSootExpr(((JIfStmt) stmt).getCondition());
                Unit nextUnit = i+1 == path.size() ? null : path.get(i + 1);
//                System.out.println("The nextUnit is: " + nextUnit);

                //如果ifstmt的后继语句不是ifstmt中goto语句，说明ifstmt中的条件为假
                if (!((JIfStmt) stmt).getTarget().equals(nextUnit))
                    ifstms = context.mkNot(ifstms);
                pathConstraint.add(ifstms);
                continue;
            }
            if (stmt instanceof JReturnStmt) {
                expectedResult = stmt.toString().replace("return", "").trim();
            }
        }
//        System.out.println("The step conditions with JimpleVars are: " + pathConstraint);
        return pathConstraint;
    }
}
